#!/usr/bin/env python3
"""
数据库设置和连接测试脚本
"""

import os
import sys
from database_manager import DatabaseManager
from database_config import DATABASE_CONFIGS

def test_database_connection(db_type: str = "sqlite"):
    """测试数据库连接"""
    print(f"=== 测试 {db_type.upper()} 数据库连接 ===")
    
    try:
        # 临时设置数据库类型
        original_type = os.environ.get("DATABASE_TYPE")
        os.environ["DATABASE_TYPE"] = db_type
        
        # 创建数据库管理器
        db = DatabaseManager()
        
        # 测试连接
        conn = db.get_connection()
        print(f"✓ 成功连接到 {db_type} 数据库")
        
        # 测试查询
        cases = db.query_cases()
        print(f"✓ 成功查询到 {len(cases)} 个案件")
        
        # 测试搜索
        search_results = db.search_cases("合同")
        print(f"✓ 搜索功能正常，找到 {len(search_results)} 个相关案件")
        
        conn.close()
        print(f"✓ {db_type.upper()} 数据库测试完成")
        
        # 恢复原始设置
        if original_type:
            os.environ["DATABASE_TYPE"] = original_type
        elif "DATABASE_TYPE" in os.environ:
            del os.environ["DATABASE_TYPE"]
            
        return True
        
    except Exception as e:
        print(f"✗ {db_type.upper()} 数据库连接失败: {e}")
        return False

def setup_database():
    """设置数据库"""
    print("=== 数据库设置向导 ===")
    print("支持的数据库类型:")
    print("1. SQLite (默认，无需配置)")
    print("2. MySQL")
    print("3. PostgreSQL")
    print("4. 自定义")
    
    choice = input("\n请选择数据库类型 (1-4): ").strip()
    
    if choice == "1" or choice == "":
        print("使用 SQLite 数据库")
        return test_database_connection("sqlite")
    
    elif choice == "2":
        print("配置 MySQL 数据库")
        print("请设置以下环境变量或编辑 .env 文件:")
        print("MYSQL_HOST=localhost")
        print("MYSQL_PORT=3306") 
        print("MYSQL_DATABASE=arbitration")
        print("MYSQL_USERNAME=root")
        print("MYSQL_PASSWORD=your_password")
        
        if input("\n已配置完成? (y/N): ").lower() == 'y':
            return test_database_connection("mysql")
    
    elif choice == "3":
        print("配置 PostgreSQL 数据库")
        print("请设置以下环境变量或编辑 .env 文件:")
        print("POSTGRES_HOST=localhost")
        print("POSTGRES_PORT=5432")
        print("POSTGRES_DATABASE=arbitration") 
        print("POSTGRES_USERNAME=postgres")
        print("POSTGRES_PASSWORD=your_password")
        
        if input("\n已配置完成? (y/N): ").lower() == 'y':
            return test_database_connection("postgresql")
    
    elif choice == "4":
        print("配置自定义数据库")
        print("请设置以下环境变量或编辑 .env 文件:")
        print("DB_TYPE=postgresql")
        print("DB_HOST=localhost")
        print("DB_PORT=5432")
        print("DB_NAME=arbitration")
        print("DB_USERNAME=your_username")
        print("DB_PASSWORD=your_password")
        
        if input("\n已配置完成? (y/N): ").lower() == 'y':
            return test_database_connection("custom")
    
    print("未进行数据库测试")
    return False

def install_database_drivers():
    """安装数据库驱动"""
    print("=== 安装数据库驱动 ===")
    print("选择要安装的驱动:")
    print("1. MySQL驱动 (mysql-connector-python)")
    print("2. PostgreSQL驱动 (psycopg2-binary)")
    print("3. 全部安装")
    
    choice = input("\n请选择 (1-3): ").strip()
    
    if choice == "1":
        os.system("pip install mysql-connector-python")
    elif choice == "2":
        os.system("pip install psycopg2-binary")
    elif choice == "3":
        os.system("pip install mysql-connector-python psycopg2-binary")
    else:
        print("未安装额外驱动")

def show_current_config():
    """显示当前数据库配置"""
    print("=== 当前数据库配置 ===")
    db_type = os.getenv("DATABASE_TYPE", "sqlite")
    config = DATABASE_CONFIGS.get(db_type, DATABASE_CONFIGS["sqlite"])
    
    print(f"数据库类型: {config['type']}")
    
    if config['type'] == 'sqlite':
        print(f"数据库文件: {config['path']}")
    else:
        for key, value in config.items():
            if key not in ['type', 'password']:
                print(f"{key}: {value}")
            elif key == 'password':
                print(f"{key}: {'***' if value else '(未设置)'}")

if __name__ == "__main__":
    print("仲裁案件数据库管理工具")
    print("=" * 30)
    
    if len(sys.argv) > 1:
        command = sys.argv[1]
        
        if command == "test":
            db_type = sys.argv[2] if len(sys.argv) > 2 else "sqlite"
            test_database_connection(db_type)
        elif command == "setup":
            setup_database()
        elif command == "install":
            install_database_drivers()
        elif command == "config":
            show_current_config()
        else:
            print("未知命令")
    else:
        print("1. 查看当前配置")
        show_current_config()
        print("\n2. 测试连接")
        test_database_connection()
        print("\n3. 如需更改配置，运行: python database_setup.py setup")
        print("4. 如需安装驱动，运行: python database_setup.py install")
